"""
网页滚动长截图工具 by 人工智能助手
支持：Chrome/Firefox/Edge | Windows/macOS/Linux
"""

from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as ec
from PIL import Image
import time
import os
import logging
from datetime import datetime


class WebScrollCapturer:
    def _merge_configs(self, global_config, task_config):
        """深度合并配置字典"""
        merged = {
            'global': global_config['global'],
            'task': task_config
        }
        # 处理浏览器配置的合并
        merged['task']['webscroll_capture']['browser'] = {
            **global_config['global']['browser'],
            **task_config['webscroll_capture'].get('browser', {})
        }
        return merged

    """
            初始化截图工具

            参数:
            driver_path : 浏览器驱动路径
            browser_type : 浏览器类型(chrome/firefox/edge)
            headless : 是否使用无头模式
            viewport_width : 视口宽度
            quality : 输出图片质量(1-100)
            temp_dir : 临时文件存储目录
            """
    def __init__(self, task_config, global_config):
        # 合并全局配置和任务配置
        self.config = self._merge_configs(global_config, task_config)

        # 初始化浏览器配置
        self.driver_path = self.config['webscroll_capture']['browser']['driver_path']
        self.service = Service(executable_path=self.config['global']['browser']['driver_path'])
        self.browser_type = self.config['task']['webscroll_capture'].get('browser_type') or self.config['global']['browser']['browser_type']
        self.headless = self.config['global']['browser']['headless']
        self.viewport_width = self.config['global']['browser']['viewport_width']
        self.quality = self.config['global']['browser']['quality']

        self.temp_dir = self.config['task']['webscroll_capture']['temp_dir']
        self.save_dir = self.config['task']['webscroll_capture']['save_dir']

        self.driver = self._init_browser()
        self.url = self.config['task']['webscroll_capture']['url']
        self.timeout = self.config['task']['webscroll_capture']['timeout']
        self.output_path = self.config['task']['webscroll_capture']['output_path']
        self.task_name = self.config['task']['webscroll_capture']['name']
        self.logger = logging.getLogger(self.task_name)

        # 创建临时目录
        os.makedirs(self.temp_dir, exist_ok=True)

    # def __exit__(self, exc_type, exc_val, exc_tb):
    #     # 增强型资源释放
    #     if self.driver:
    #         self.driver.quit()
    #     if self.service:
    #         self.service.stop()
    #     # os.system("taskkill /f /im chromedriver.exe")  # 确保进程终止

    def _init_browser(self):
        """初始化浏览器实例"""
        options = None

        if self.browser_type == 'chrome':
            chrome_testing_path = r'F:\pycode\chrome-win64\chrome-win64\chrome.exe'
            options = webdriver.ChromeOptions()
            options.binary_location = chrome_testing_path
            options.add_experimental_option('detach', True)
            options.add_argument(f'--window-size={self.viewport_width},1080')
            if self.headless:
                options.add_argument('--headless=new')
            options.add_argument('--disable-gpu')
            options.add_argument('--no-sandbox')
            options.add_argument('--disable-dev-shm-usage')

            return webdriver.Chrome(
                service=self.service,
                options=options
            )

        elif self.browser_type == 'firefox':
            options = webdriver.FirefoxOptions()
            if self.headless:
                options.add_argument('-headless')
            options.add_argument('--width={}'.format(self.viewport_width))
            options.add_argument('--height=1080')
            return webdriver.Firefox(
                options=options
            )

        else:
            raise ValueError("不支持的浏览器类型")

    def _smart_scroll(self, scroll_delay: float = 0.5, max_attempts: int = 3):
        """
        智能滚动加载页面内容
        """
        attempts = 0
        last_height = self.driver.execute_script("return document.body.scrollHeight")

        while attempts < max_attempts:
            # 滚动到底部
            self.driver.execute_script("window.scrollTo(0, document.body.scrollHeight);")
            time.sleep(scroll_delay)

            # 检测新高度
            new_height = self.driver.execute_script("return document.body.scrollHeight")
            if new_height == last_height:
                attempts += 1
            else:
                attempts = 0
                last_height = new_height

    def _remove_interfering_elements(self):
        """移除干扰元素"""
        self.driver.execute_script("""
            const selectors = [
                'header', 'footer', '.fixed-top', 
                '.cookie-banner', '.ad-container',
                'iframe', '.popup'
            ];

            // 隐藏固定定位元素
            document.querySelectorAll(selectors.join(','))
                .forEach(el => el.style.display = 'none');

            // 禁用动画效果
            const style = document.createElement('style');
            style.textContent = `
                * {
                    animation: none !important;
                    transition: none !important;
                }
            `;
            document.head.appendChild(style);
        """)

    def force_full_render(self):
        """强制完整渲染页面所有元素"""
        # 禁用可能影响布局的动态效果
        self.driver.execute_script("""
            document.body.style.overflow = 'visible';
            document.body.style.transform = 'none';
            const all = document.querySelectorAll('*');
            all.forEach(el => {
                el.style.transition = 'none !important';
                el.style.animation = 'none !important';
            });
        """)
        time.sleep(1)  # 等待样式生效

    def get_full_dimensions(self):
        """获取页面真实尺寸（兼容横向/纵向溢出）"""
        return self.driver.execute_script("""
            const body = document.body;
            const html = document.documentElement;

            // 横向/纵向综合计算
            const width = Math.max(
                body.scrollWidth, body.offsetWidth,
                html.clientWidth, html.scrollWidth, html.offsetWidth,
                document.getElementById('report1_reportDiv').scrollWidth
            );

            const height = Math.max(
                body.scrollHeight, body.offsetHeight,
                html.clientHeight, html.scrollHeight, html.offsetHeight,
                document.getElementById('report1_reportDiv').scrollHeight
            );

            return [width, height];
        """)


    def capture(self):
        """
        执行截图操作

        参数:
        url : 目标网页URL
        output_path : 输出图片路径
        timeout : 页面加载超时时间(秒)
        """
        self.logger.info("start webscroll capture task ...")
        try:
            # 访问目标页面
            self.driver.get(self.url)

            # 等待页面基础加载完成
            WebDriverWait(self.driver, self.timeout).until(
                ec.presence_of_element_located((By.TAG_NAME, 'body'))
            )

            # 智能滚动加载内容
            self._smart_scroll()

            # 移除干扰元素
            self._remove_interfering_elements()

            # 获取最终页面尺寸
            # 等待核心内容加载（自定义关键元素）
            WebDriverWait(self.driver, self.timeout).until(
                ec.presence_of_all_elements_located((By.CSS_SELECTOR, "#report1_reportDiv table tr"))
            )

            # 强制完整渲染
            self.force_full_render()

            # 获取真实页面尺寸
            width, height = self.get_full_dimensions()
            print(f"页面真实尺寸: {width}x{height}px")

            # 设置浏览器视口（增加 50px 防止边缘裁剪）
            # self.driver.set_window_size(width + 50, height + 50)

            self.driver.set_window_size(width+100, height+300)
            # 调整浏览器窗口尺寸
            # self.driver.set_window_size(total_width, total_height)
            time.sleep(1)  # 等待渲染稳定

            # 执行截图并拼接
            self._stitch_screenshot(self.output_path)

        except Exception as e:
            self.logger.error(f"capture failed: {str(e)}")
            raise RuntimeError(f"截图失败: {str(e)}")
        # finally:
        #     self.logger.info("stop webscroll capture task ...")
        #     self.driver.quit()
        #     self.service.stop()
            # 清理临时文件
            # for f in os.listdir(self.temp_dir):
            #     os.remove(os.path.join(self.temp_dir, f))
            # os.rmdir(self.temp_dir)

    def _stitch_screenshot(self, output_path: str):
        """分块截图并拼接"""
        viewport_height = self.driver.execute_script("return window.innerHeight")
        total_height = self.driver.execute_script("return document.body.scrollHeight")
        # viewport_height = 1080
        # total_height = 2500
        temp_files = []
        current_position = 0

        try:
            # 分块截图
            while current_position < total_height:
                # 滚动到指定位置
                self.driver.execute_script(f"window.scrollTo(0, {current_position})")
                time.sleep(0.3)  # 等待渲染

                # 保存临时文件
                temp_file = os.path.join(self.temp_dir, f"part_{current_position}.png")
                self.driver.save_screenshot(temp_file)
                temp_files.append(temp_file)

                current_position += viewport_height

            # 使用PIL拼接图片
            images = [Image.open(f) for f in temp_files]
            total_width = images[0].width
            combined_image = Image.new(
                'RGB',
                (total_width, total_height),
                color=(255, 255, 255)
            )

            y_offset = 0
            for img in images:
                # 计算实际裁剪区域
                crop_height = min(img.height, total_height - y_offset)
                combined_image.paste(
                    img.crop((0, 0, total_width, crop_height)),
                    (0, y_offset)
                )
                y_offset += crop_height

            # 保存最终结果
            # save_path = os.path.join(save_dir, f"screenshot_{timestamp}.png")
            # 每天覆盖原来的图片
            save_path = os.path.join(self.save_dir, output_path)
            combined_image.save(save_path, quality=self.quality)
            # print(f"✅ 截图成功保存至: {os.path.abspath(output_path)}")
            print(f"✅ 截图成功保存至: {save_path}")
            self.logger.info(f"[Success] {datetime.now().strftime('%H:%M:%S')} save image:{save_path}")

        except Exception as e:
            raise RuntimeError(f"图片拼接失败: {str(e)}")



if __name__ == "__main__":
    # 使用示例
    TARGET_URL = "https://10.200.95.111:28080/runqianReport/report/jsp/dhccpmrunqianreport.jsp?report=__u6DF1__u5733__u5E02__u5B9D__u5B89__u533A__u677E__u5C97__u4EBA__u6C11__u533B__u9662__u7B80__u62A5.rpx&reportUrlType=1&cspsessionid=Ws2ov7DSEB&LayoutManager=cn_iptcp%3A10.200.95.111%5B51773%5D%3ADHC-APP&ServerSideRedirect=dhccpmrunqianreport.csp&TMENU=61285&TPAGID=1225388125&a=a&reportName=__u6DF1__u5733__u5E02__u5B9D__u5B89__u533A__u677E__u5C97__u4EBA__u6C11__u533B__u9662__u7B80__u62A5.rpx"  # 目标网页地址
    SAVE_DIR = "./screenshots"  # 截图保存目录
    config = {
        'driver_path': r'F:\pycode\ReportRobot\chromedriver-win64\chromedriver-win64\chromedriver.exe',  # 驱动路径
        'browser_type': 'chrome',  # 浏览器类型
        'headless': True,  # 无头模式
        'viewport_width': 1440,  # 视口宽度
        'quality': 95  # 图片质量
    }

    try:
        capturer = WebScrollCapturer(**config)
        capturer.capture()
    except Exception as e:
        print(f"❌ 错误发生: {str(e)}")